import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from model import ImageClassifier
from dataset import get_cifar10_data_loaders
from train import train_model, evaluate_model
from utils import save_model, load_model
import matplotlib.pyplot as plt

# CIFAR10类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

def main():
    """
    主函数，执行模型训练和评估流程
    """
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    # 设置随机种子以保证可重复性
    torch.manual_seed(42)
    if device == 'cuda':
        torch.cuda.manual_seed_all(42)
    
    # 获取数据加载器
    batch_size = 128  # 增大批次大小
    train_loader, test_loader = get_cifar10_data_loaders(batch_size=batch_size)
    
    # 初始化模型、损失函数和优化器
    model = ImageClassifier().to(device)
    
    # 使用交叉熵损失
    criterion = nn.CrossEntropyLoss()
    
    # 使用AdamW优化器(带权重衰减)
    optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-4)
    
    # 训练模型
    num_epochs = 100  # 增加训练轮数
    print("Starting training...")
    trained_model, train_losses, train_accuracies = train_model(
        model, train_loader, criterion, optimizer, device, num_epochs
    )
    
    # 评估模型
    # ... (前面的代码保持不变)
    print("Evaluating model...")
    accuracy, test_loss, class_accuracies = evaluate_model(
        trained_model, 
        test_loader, 
        device,
        classes=classes  # 传递classes参数
    )
    print(f"Final Test Accuracy: {accuracy:.2f}%")

# ... (后面的代码保持不变)
    
    # 保存模型
    save_model(trained_model, "improved_cifar10_classifier.pth")
    print("Model saved to improved_cifar10_classifier.pth")
    
    # 绘制类别准确率柱状图
    plt.figure(figsize=(10, 5))
    plt.bar(range(10), class_accuracies)
    plt.xticks(range(10), classes, rotation=45)
    plt.xlabel('Class')
    plt.ylabel('Accuracy (%)')
    plt.title('Class-wise Accuracy')
    plt.tight_layout()
    plt.savefig('class_accuracy.png')
    plt.close()

if __name__ == "__main__":
    main()